#import matplotlib.font_manager as fm
#myfont = fm.FontProperties(fname='/usr/share/fonts/font/simhei.ttf')
#myfont = fm.FontProperties(fname='/usr/share/fonts/font/msyhbd.ttc')

import mplhep
import matplotlib.pyplot as plt
plt.style.use(mplhep.style.LHCb2)

import pandas as pd
import numpy as np

class colors_gen(object):
    def __init__(self):
        self.__colors = [
            "red",
            "orange",
            "purple",
            "springgreen",
            "y",
            "blue",
            "c",
            "m",
        ]
        self.__i_color = 0

    def a_color(self):
        i = self.__i_color
        if i >= len(self.__colors):
            self.__i_color = i = 0
        self.__i_color = self.__i_color + 1
        return self.__colors[i]

def plot_a_year(fig, iyear, coordinates):
    plt.clf()
    plt.xlim(0, 360.)
    plt.ylim(-90., 90.)
    #print(coordinates)

    ax = fig.get_axes()[0]
    #ax.set_xlabel("赤经（度）", fontproperties=myfont, fontsize=30)
    #ax.set_ylabel("赤纬（度）", fontproperties=myfont, fontsize=30)
    ax.set_xlabel("Right ascension / deg")
    ax.set_ylabel("Declination / deg")

    last_point = coordinates[0]

    # 标记年份
    plt.text(50., 65., s=f"Year {iyear + 1}", fontsize=30, fontweight='bold', ha="right", va="bottom", color='black')
    #plt.text(50., 65., s=f"第{iyear + 1}年", fontsize=30, fontweight='bold', ha="right", va="bottom", color='black', fontproperties=myfont)

    # 标记起始点
    plt.annotate(text=f"{int(last_point[3])}", xy=(last_point[0], last_point[1]), xytext=(last_point[0]+4, last_point[1]-5), color='green')
    plt.scatter(last_point[0], last_point[1], marker='*', color='green')
    if len(coordinates) == 1:
        return

    colors = colors_gen()

    for ipoint in coordinates[1:]:
        plt.plot ([last_point[0], ipoint[0]], [last_point[1], ipoint[1]], linewidth=2, color='black', zorder=-2)
        plt.arrow(last_point[0], last_point[1], \
                  (ipoint[0] - last_point[0])/2., (ipoint[1] - last_point[1])/2., \
                  width = 1., linewidth=0.1, color='black', zorder=-1)
        
        a_color = colors.a_color()
        plt.scatter(ipoint[0], ipoint[1], s=30, marker='.', color=a_color)
        plt.annotate(text=f"{int(ipoint[3])}", xy=(ipoint[0], ipoint[1]), xytext=(ipoint[0]+2, ipoint[1]-5), color=a_color)
        #plt.text(ipoint[0], ipoint[1], s=f"{int(ipoint[3])}", ha="right", va="top", color=a_color)
        #print(ipoint)
        last_point = ipoint

    #plt.scatter(coordinates[-1][0], coordinates[-1][1], marker='D', color='red')


def plot_mission_route(file_path, save_path, end_time, param_list):

    fig = plt.figure(figsize=[12,6])

    df = pd.read_csv(file_path)
    coordinate_list = np.array(df[param_list + ['Time', 'ID']])    

    last_year = 1
    coordinate_a_year = []
    coordinate_years = []
    for ipoint in coordinate_list:
        if end_time > 0 and ipoint[2] > end_time:
            break
        if np.floor(ipoint[2] / 365.2242) >= last_year:
            last_year = last_year + 1
            coordinate_years.append(coordinate_a_year)
            coordinate_a_year = []
        coordinate_a_year.append(ipoint.tolist())
    coordinate_years.append(coordinate_a_year)

    print(f"The mission needs {len(coordinate_years)} years in total.")
    for i, ayear in enumerate(coordinate_years):
        if len(ayear) == 0:
            continue
        print(f"Plot year {i+1}. Task num: {len(ayear)}.")
        plot_a_year(fig=fig, iyear=i, coordinates=ayear)
        plt.savefig(f"{save_path}/year_{i+1}.png")

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Test Operation module")
    #parser.add_argument("-f", "--file", default="result.csv", dest="file")
    parser.add_argument(dest="file", type=str)
    parser.add_argument("-s", "--save_path", default="default", dest="save_path")
    parser.add_argument("-t", "--end_time", type=float, default=-1, dest="end_time")
    arglist = parser.parse_args()
    
    import os
    csv_file = arglist.file
    if not os.path.exists(csv_file):
        raise FileNotFoundError(f"No such file or directory: '{csv_file}'.")
    save_path = csv_file[:csv_file.rfind("/")] + "/mission_route"
    if arglist.save_path != "default":
        save_path = arglist.save_path
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    plot_mission_route(file_path=arglist.file, save_path=save_path, end_time=arglist.end_time, param_list=['ra', 'dec'])

if __name__ == "__main__":
    main()